#coding: utf-8
from mrjob.job import MRJob
from mrjob.step import MRStep
import math


def ruzicka_reducer(uni1, uni2, cross_values):
    '''
    输入：
        uni1: 集合1的Uni值
        uni2: 集合2的Uni值
        cross_values: (f_{i, k}, f{j, k})的iterable对象
    输出:
        计算 conj = sum(min(f_{i, k}, f{j, k}))
        输出 conj / (uni1 + uni2 - conj)
    '''
    if uni1 == 0 or uni2 == 0:
        return 0.
    conj = 0.
    for x, y in cross_values:
        conj += min(x, y)
    return float(conj) / (uni1 + uni2 - conj)


def cos_reducer(uni1, uni2, cross_values):
    '''
    输入：
        uni1: 集合1的Uni值
        uni2: 集合2的Uni值
        cross_values: (f_{i, k}, f{j, k})的iterable对象
    输出:
        计算 conj = sum(min(f_{i, k}, f{j, k}))
        输出 conj / sqrt(uni1 * uni2)
    '''
    if uni1 == 0 or uni2 == 0:
        return 0.
    conj = 0.
    for x, y in cross_values:
        conj += min(x, y)
    return float(conj) / math.sqrt(uni1 * uni2)


def dice_reducer(uni1, uni2, cross_values):
    '''
    输入：
        uni1: 集合1的Uni值
        uni2: 集合2的Uni值
        cross_values: (f_{i, k}, f{j, k})的iterable对象
    输出:
        计算 conj = sum(min(f_{i, k}, f{j, k}))
        输出 2 * (conj / (uni1 + uni2))
    '''
    if uni1 == 0 or uni2 == 0:
        return 0.
    conj = 0.
    for x, y in cross_values:
        conj += min(x, y)
    return 2 * float(conj) / (uni1 + uni2)


class VSmartSim(MRJob):

    def configure_args(self):
        super(VSmartSim, self).configure_args()
        self.add_passthru_arg(
            '--sim-type', default='ruzicka', choices=['ruzicka_reducer', 'cos', 'dice'],
            help="Specify the type of similarity")
        self.add_passthru_arg(
            '--q', default=10000, type=int)

    def mapper_joining_oa(self, _, line):
        '''
        Online-Aggregation Joining mapper
        输入：
            M_i, a_k, f_{i, k}
        输出：
            M_i, (0, f_{i, k})
            M_i, (1, (a_k, f_{i, k}))
        '''
        line = line.split(',')
        M_i = int(line[0])
        a_k = line[1]
        f_ik = int(line[2])
        yield M_i, (0, f_ik)
        yield M_i, (1, (a_k, f_ik))

    def reducer_joining_oa(self, M_i, values):
        '''
        Online-Aggregation Joining reducer
        输入：
            M_i, (0, f_{i, k})
            M_i, (1, (a_k, f_{i, k}))
        输出：
            None, (M_i, Uni(M_i), (a_k, f_{i, k}))

        没有secondary key的方法，全部读到内存中
        '''
        f = []
        m = []
        for value in values:
            if value[0] == 0:
                f.append(value[1])
            else:
                m.append(value[1])
        uni = sum(f)
        for pair in m:
            yield None, (M_i, uni, pair)

    def mapper_sim1(self, _, value):
        '''
        Similarity Stage1 mapper
        输入：
            None, (M_i, Uni(M_i), (a_k, f_{i, k}))
        输出：
            a_k, (M_i, Uni(M_i), f_{i, k})
        '''
        M_i, uni, (a_k, f_ik) = value
        yield a_k, (M_i, uni, f_ik)

    def reducer_sim1(self, key, values):
        '''
        Similarity Stage1 reducer
        输入：
            a_k, (M_i, Uni(M_i), f_{i, k})
        输出：
            (M_i, M_j, Uni(M_i), Uni(M_j)), (f_{i, k}, f_{j, k})
        values的长度大于一个数值（论文中的q）会被全部舍弃掉。
        '''
        tmp_values = []
        for i, value in enumerate(values):
            if i + 1 > self.options.q:
                return
            tmp_values.append(value)
        values = tmp_values
        length = len(values)
        for i in range(length - 1):
            for j in range(i + 1, length):
                M_i, uni_i, f_ik = values[i]
                M_j, uni_j, f_jk = values[j]
                yield (M_i, M_j, uni_i, uni_j), (f_ik, f_jk)

    def mapper_sim2(self, key, value):
        '''
        Similarity Stage2 mapper

        什么都不做。
        '''
        yield key, value

    def reducer_sim2(self, key, values):
        '''
        Similarity Stage2 reducer

        输入：
            (M_i, M_j, Uni(M_i), Uni(M_j)), (f_{i, k}, f_{j, k})
        输出：
            (M_i, M_j), Sim(M_i, M_j)
        '''
        if self.options.sim_type == "ruzicka":
            sim = ruzicka_reducer(key[2], key[3], values)
        elif self.options.sim_type == "cos":
            sim = cos_reducer(key[2], key[3], values)
        elif self.options.sim_type == "dice":
            sim = dice_reducer(key[2], key[3], values)
        else:
            sim = None
        yield (key[0], key[1]), sim

    def steps(self):
        return [
            MRStep(
                mapper=self.mapper_joining_oa,
                reducer=self.reducer_joining_oa
            ),
            MRStep(
                mapper=self.mapper_sim1,
                reducer=self.reducer_sim1,
            ),
            MRStep(
                mapper=self.mapper_sim2,
                reducer=self.reducer_sim2,
            ),
        ]


if __name__ == "__main__":
    VSmartSim.run()
